# 单独运行，根据标签报文切分
import json
import dpkt
import os

round_cnt = 1
fd = open("config.json", "r")
config = json.load(fd)
fd.close()
initial_path = config["initial_store_path"]
label_path = config["label_store_path"]


def main():
    global round_cnt
    for pcap_file in os.listdir(initial_path):
        print("[INFO] Processing {}".format(pcap_file))
        process_signal_file(os.path.join(initial_path, pcap_file), label_path)
        round_cnt += 1


def process_signal_file(file_path, label_path):
    global round_cnt
    fd = open(file_path, "rb")
    magic_head = fd.read(4)
    fd.seek(0, 0)
    pcap_reader = None
    if magic_head == b'\n\r\r\n':
        pcap_reader = dpkt.pcapng.Reader(fd)
    elif magic_head == b'\xd4\xc3\xb2\xa1':
        pcap_reader = dpkt.pcap.Reader(fd)
    else:
        print("[DEBUG in PcapUtils] It is not a pcap or pcapng file")
    new_page_flag = False
    writer_fd = None
    pcap_writer = None
    page_cnt = 1
    for ts, buf in pcap_reader.readpkts():
        # print(ts)
        eth = dpkt.ethernet.Ethernet(buf)
        if eth.type != dpkt.ethernet.ETH_TYPE_IP:
            if pcap_writer is not None:
                pcap_writer.writepkt(buf, ts)
            continue
        ip = eth.data
        if ip.dst != b'\x7b\x7b\x7b\x7b':
            if pcap_writer is not None:
                pcap_writer.writepkt(buf, ts)
            continue
        new_page_flag = not new_page_flag
        if new_page_flag:
            store_path = "{}\\page{}".format(label_path, page_cnt)
            if not os.path.exists(store_path):
                os.mkdir(store_path)
            writer_fd = open("{}\\page{}_r{}.pcapng".format(store_path, page_cnt, round_cnt), "wb")
            pcap_writer = dpkt.pcapng.Writer(writer_fd)
            pcap_writer.writepkt(buf, ts)
            page_cnt += 1
        elif pcap_writer is not None:
            pcap_writer.writepkt(buf, ts)
            pcap_writer.close()
            writer_fd.close()
            pcap_writer = None
        else:
            continue


if __name__ == "__main__":
    main()

